package cz.cuni.amis.episodic.lisp.netcreators;

import java.util.List;
import java.util.Set;

import lombok.AllArgsConstructor;
import smile.Network;

/**
 * Sequence of altering Goal and Hidden nodes. Hidden nodes are learned by EM.
 *
 * @author ik
 */

public class HHMMCreator extends NetCreator {

    int hiddenNodeStatesNum;

    boolean observationTemporalModel;
    
    public HHMMCreator(int hiddenNodeStatesNum, boolean observationTemporalModel) {
        this.hiddenNodeStatesNum = hiddenNodeStatesNum;
        filenameSuffix = "_hhmm";
        this.observationTemporalModel = observationTemporalModel;
    }

    
    
    @Override
    public Network createNet(List<List<String>> levelActions) {
        final int max = levelActions.size();

        Network net = new Network();
        String[] nodes = createNodeNames(max);
        for (int i = 0; i < max - 1; i++) {
            int id = max - i - 1;
            int nodeIx = 2 * i;

            // G
            addNode(net, nodes[nodeIx], levelActions.get(i));
            // H
            addNode(net, nodes[nodeIx + 1], hiddenNodeStatesNum);
            if (i > 0) { // H_(i-1) -> G_i
                net.addArc(nodes[nodeIx - 1], nodes[nodeIx]);
            }
            // G_i -> H_i
            net.addArc(nodes[nodeIx], nodes[nodeIx + 1]);
        }
        // observation O
        addNode(net, nodes[nodes.length - 1], levelActions.get(levelActions.size() - 1));
        // H_last -> O
        net.addArc(nodes[nodes.length - 2], nodes[nodes.length - 1]);

        // add temporal arcs
        adjustNodes(net, nodes, observationTemporalModel);

        return net;

    }

    public static void addNode(Network net, String name, int statesNum) {
        addTemporalNode(net, name);
        for (int i = 0; i < statesNum; i++) {
            net.addOutcome(name, "HiddenState" + i);
        }
    }
    
    
    public static String[] createNodeNames(int levels) {
        String[] nodes = new String[levels * 2 - 1];
        for (int i = 0; i < levels - 1; i++) {
            int id = levels - i - 1;
            int nodeIx = 2 * i;
            // goal node
            nodes[nodeIx] = "G" + id;
            // hidden node representing goal's internal structure
            nodes[nodeIx + 1] = "H" + id;
        }
        nodes[nodes.length - 1] = "O";
        return nodes;
    }
}
